import cv2
import numpy as np
from utils import utils, ocr_utils


def table_split(original_image):
    a = 80
    b = 30

    # 先对图像规定尺寸
    bgr_image = original_image.copy()
    # 获取图像特征
    height, width, _ = original_image.shape
    # 灰度图
    gray_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2GRAY)
    # 二值化
    thresh_image = cv2.threshold(gray_image, 0, 255, cv2.THRESH_OTSU)[1]
    # 闭运算
    kernel = np.ones((a, b), np.uint8)  # 这个对准确度关系大
    closing = cv2.erode(thresh_image, kernel, iterations=1)
    # 轮廓检测
    contours, hierarchy = cv2.findContours(closing, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    # 存放检测到的所有表格
    table_list = []
    # 遍历所有轮廓
    for contour in contours:
        # 计算轮廓周长和面积
        perimeter = cv2.arcLength(contour, True)
        area = cv2.contourArea(contour, True)
        # 进行多边形逼近
        approx = cv2.approxPolyDP(contour, 0.08 * perimeter, True)
        # 如果轮廓拥有四个顶点，且面积达到一定值，则认为其是一个表格
        if len(approx) == 4 and area > width * height * 0.01:
            approx = utils.approx_sort(approx)
            # 划线绿色线条
            cv2.drawContours(bgr_image, [approx], 0, (0, 255, 0), 10)
            # 利用approx的轮廓对table进行透视变换
            pts1 = np.float32(approx)
            pts2 = np.float32([[0, 0], [width, 0], [width, height], [0, height]])
            # 获取透视变换矩阵
            M = cv2.getPerspectiveTransform(pts1, pts2)
            # 进行透视变换
            table = cv2.warpPerspective(original_image, M, (width, height))
            ocr_utils.cv_show("table", table)
            table_list.append(table)
    return table_list
